package fem2;

import fem2.enu.MassMatrixType;
import inf.jlinalg.Array2DMatrix;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class ParallelAssembler extends Assembler {

	class KtCompute implements Runnable {

		private Element e;
		private IMatrix kt;
		private ExecutorService ex;;

		public KtCompute(Element e, IMatrix kt, ExecutorService ex) {
			this.e = e;
			this.kt = kt;
			this.ex = ex;
		}

		@Override
		public void run() {
			int n = e.getMeshPart().countDofs();
			IMatrix kte = new Array2DMatrix(n, n);
			kte = (IMatrix) e.computeKte(kte);
			ex.execute(new KtScatter(e, kt, kte));
		}

	}

	class KtScatter implements Runnable {

		private Element e;
		private IMatrix kt;
		private IMatrix kte;

		public KtScatter(Element e, IMatrix kt, IMatrix kte) {
			this.e = e;
			this.kt = kt;
			this.kte = kte;
		}

		@Override
		public void run() {
			int[] dofs = e.getMeshPart().getDofNumbers();

			for (int j = 0; j < kte.getRowCount(); j++) {
				for (int k = 0; k < kte.getColumnCount(); k++) {
					if (dofs[j] != -1 && dofs[k] != -1) {
						kt.add(dofs[j], dofs[k], kte.get(j, k));
					}
				}
			}

		}
	}

	class RiCompute implements Runnable {

		private Element e;
		private double[] ri;
		private ExecutorService ex;;

		public RiCompute(Element e, double[] ri, ExecutorService ex) {
			this.e = e;
			this.ri = ri;
			this.ex = ex;
		}

		@Override
		public void run() {
			int n = e.getMeshPart().countDofs();
			IVector rie = new ArrayVector(n);
			rie = (IVector) e.computeRie(rie);
			ex.execute(new RiScatter(e, ri, rie));
		}

	}

	class RiScatter implements Runnable {

		private Element e;
		private double[] ri;
		private IVector rie;

		public RiScatter(Element e, double[] ri, IVector rie) {
			this.e = e;
			this.ri = ri;
			this.rie = rie;
		}

		@Override
		public void run() {
			int[] dofs = e.getMeshPart().getDofNumbers();

			for (int j = 0; j < dofs.length; j++) {
				if (dofs[j] != -1) {
					ri[dofs[j]] += rie.get(j);
				}
			}

		}
	}

	class MCompute implements Runnable {

		private Element e;
		private IMatrix m;
		private MassMatrixType massType;
		private ExecutorService ex;;

		public MCompute(Element e, IMatrix m, MassMatrixType massType, ExecutorService ex) {
			this.e = e;
			this.m = m;
			this.massType = massType;
			this.ex = ex;
		}

		@Override
		public void run() {
			int n = e.getMeshPart().countDofs();
			IMatrix me = new Array2DMatrix(n, n);
			me = (IMatrix) e.computeMe(me, massType);
			ex.execute(new MScatter(e, m, me));
		}

	}

	class MScatter implements Runnable {

		private Element e;
		private IMatrix m;
		private IMatrix me;

		public MScatter(Element e, IMatrix m, IMatrix me) {
			this.e = e;
			this.m = m;
			this.me = me;
		}

		@Override
		public void run() {
			int[] dofs = e.getMeshPart().getDofNumbers();

			for (int j = 0; j < me.getRowCount(); j++) {
				for (int k = 0; k < me.getColumnCount(); k++) {
					if (dofs[j] != -1 && dofs[k] != -1) {
						m.add(dofs[j], dofs[k], me.get(j, k));
					}
				}
			}

		}
	}

	private ExecutorService executor1;
	private ExecutorService executor2;
	private int numThreads;

	public ParallelAssembler(Model m, int numThreads) {
		super(m);
		this.numThreads = numThreads;
	}

	@Override
	public void assembleKt(IMatrix kt) {
		executor1 = Executors.newFixedThreadPool(numThreads);
		executor2 = Executors.newSingleThreadExecutor();

		Model m = getModel();
		for (int i = 0; i < m.countElements(); i++) {
			Element e = m.getElement(i);
			if (!e.isInactive())
				executor1.execute(new KtCompute(e, kt, this.executor2));
		}
		executor1.shutdown();
		try {
			executor1.awaitTermination(10, TimeUnit.HOURS);
		} catch (InterruptedException e) {
			throw new Error();
		}
		executor2.shutdown();
		try {
			executor2.awaitTermination(10, TimeUnit.HOURS);
		} catch (InterruptedException e) {
			throw new Error();
		}
	}

	@Override
	public void assembleRi(double[] ri) {
		executor1 = Executors.newFixedThreadPool(numThreads);
		executor2 = Executors.newSingleThreadExecutor();

		Model m = getModel();
		for (int i = 0; i < m.countElements(); i++) {
			Element e = m.getElement(i);
			if (!e.isInactive())
				executor1.execute(new RiCompute(e, ri, this.executor2));
		}
		executor1.shutdown();
		try {
			executor1.awaitTermination(10, TimeUnit.HOURS);
		} catch (InterruptedException e) {
			throw new Error();
		}
		executor2.shutdown();
		try {
			executor2.awaitTermination(10, TimeUnit.HOURS);
		} catch (InterruptedException e) {
			throw new Error();
		}

	}

	@Override
	public void assembleM(IMatrix m, MassMatrixType massType) {
		executor1 = Executors.newFixedThreadPool(numThreads);
		executor2 = Executors.newSingleThreadExecutor();

		Model M = getModel();
		for (int i = 0; i < M.countElements(); i++) {
			Element e = M.getElement(i);
			if (!e.isInactive())
				executor1.execute(new MCompute(e, m, massType, this.executor2));
		}
		executor1.shutdown();
		try {
			executor1.awaitTermination(10, TimeUnit.HOURS);
		} catch (InterruptedException e) {
			throw new Error();
		}
		executor2.shutdown();
		try {
			executor2.awaitTermination(10, TimeUnit.HOURS);
		} catch (InterruptedException e) {
			throw new Error();
		}
	}

}
